import cv2
import supervision as sv
import matplotlib.pyplot as plt
import time
import os
import numpy as np
from PIL import Image, ImageDraw, ImageFont
try:
    from .config_manager import config_manager
except ImportError:
    from config_manager import config_manager

# 确保中文显示正常
plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "Arial Unicode MS", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

class Visualizer:
    """结果可视化模块，负责检测结果的展示和保存"""
    
    def __init__(self, results_dir="results"):
        """初始化可视化工具
        
        Args:
            results_dir: 结果保存目录
        """
        # 初始化可视化工具
        self.box_annotator = sv.BoxAnnotator(
            thickness=2
        )
        self.label_annotator = sv.LabelAnnotator()
        
        # 创建结果目录
        self.results_dir = results_dir
        if not os.path.exists(self.results_dir):
            os.makedirs(self.results_dir)
        
        # 颜色配置
        self.colors = {
            "person": (0, 255, 0),    # 绿色
            "car": (255, 0, 0),      # 蓝色
            "truck": (0, 0, 255),    # 红色
            "bus": (255, 255, 0),    # 黄色
            "motorcycle": (255, 0, 255), # 紫色
            "bicycle": (0, 255, 255), # 青色
            "default": (128, 128, 128)  # 灰色
        }
        
        # 初始化中文字体
        self.font = self._load_chinese_font()
    
    def _load_chinese_font(self):
        """加载中文字体
        
        Returns:
            PIL字体对象，如果加载失败则返回None
        """
        try:
            # 从配置文件获取字体路径
            font_paths = config_manager.get_font_paths()
            
            # 尝试加载系统中文字体
            for font_path in font_paths:
                if os.path.exists(font_path):
                    return ImageFont.truetype(font_path, 16)
            
            # 如果都找不到，使用默认字体
            return ImageFont.load_default()
        except Exception as e:
            print(f"加载中文字体失败: {e}")
            return ImageFont.load_default()
    
    def _put_chinese_text(self, img, text, position, font_size=16, color=(0, 255, 0)):
        """在图像上添加中文文本
        
        Args:
            img: OpenCV图像 (BGR格式)
            text: 要显示的文本
            position: 文本位置 (x, y)
            font_size: 字体大小
            color: 文本颜色 (BGR格式)
        
        Returns:
            添加文本后的图像
        """
        try:
            # 将OpenCV图像转换为PIL图像
            img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(img_pil)
            
            # 尝试加载字体（多重回退机制）
            font = None
            font_paths = config_manager.get_font_paths()
            
            for font_path in font_paths:
                try:
                    if os.path.exists(font_path):
                        font = ImageFont.truetype(font_path, font_size)
                        break
                except:
                    continue
            
            # 如果所有字体都加载失败，使用默认字体
            if font is None:
                font = self.font or ImageFont.load_default()
            
            # 转换颜色格式 (BGR -> RGB)
            color_rgb = (color[2], color[1], color[0])
            
            # 绘制文本
            draw.text(position, text, font=font, fill=color_rgb)
            
            # 转换回OpenCV格式
            img_cv = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
            return img_cv
        except Exception as e:
            print(f"PIL中文文本显示失败: {e}")
            # 如果PIL方法完全失败，尝试使用ASCII字符替代
            try:
                # 将中文字符转换为拼音或英文描述（简单处理）
                ascii_text = self._convert_to_ascii(text)
                cv2.putText(img, ascii_text, position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
            except:
                # 最后的回退：显示占位符
                cv2.putText(img, "[Chinese Text]", position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
            return img
    
    def _convert_to_ascii(self, text):
        """将中文文本转换为ASCII可显示的文本
        
        Args:
            text: 输入文本
        
        Returns:
            ASCII文本
        """
        # 使用配置管理器进行文本转换
        ascii_text = config_manager.convert_to_ascii(text)
        
        # 移除剩余的非ASCII字符
        ascii_text = ''.join(char if ord(char) < 128 else '?' for char in ascii_text)
        
        return ascii_text
    
    def visualize_results(self, frame, detections, labels=None, description=""):
        """使用Supervision进行结果可视化
        
        Args:
            frame: 输入图像
            detections: 检测结果
            labels: 标签列表
            description: 场景描述
        
        Returns:
            可视化后的图像
        """
        # 绘制边界框
        annotated_frame = self.box_annotator.annotate(
            scene=frame.copy(),
            detections=detections
        )
        
        # 绘制标签（如果提供了标签）
        if labels is not None:
            annotated_frame = self.label_annotator.annotate(
                scene=annotated_frame,
                detections=detections,
                labels=labels
            )
        
        # 在图像上添加场景描述（使用中文字体支持）
        if description:
            y_offset = 30
            for line in description.split("\n"):
                if line.strip():  # 确保不是空行
                    # 使用新的中文文本显示方法
                    annotated_frame = self._put_chinese_text(
                        annotated_frame,
                        line.strip(),
                        (10, y_offset),
                        font_size=14,
                        color=(0, 255, 0)  # 绿色
                    )
                    y_offset += 22
        
        # 添加时间戳（英文数字可以正常显示）
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        cv2.putText(
            annotated_frame,
            timestamp,
            (annotated_frame.shape[1] - 180, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (255, 0, 0),
            1
        )
        
        return annotated_frame
    
    def save_results(self, frame, filename="result", is_rgb=True):
        """保存分析结果
        
        Args:
            frame: 要保存的图像
            filename: 文件名前缀
            is_rgb: 指示输入图像是否为RGB格式
        
        Returns:
            保存路径
        """
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filepath = os.path.join(self.results_dir, f"{filename}_{timestamp}.jpg")
        
        # 根据输入格式正确保存图像（OpenCV的imwrite需要BGR格式）
        if len(frame.shape) == 3 and frame.shape[2] == 3:
            if is_rgb:
                # 如果输入是RGB格式，转换为BGR再保存
                cv2.imwrite(filepath, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
            else:
                # 如果输入已经是BGR格式，直接保存
                cv2.imwrite(filepath, frame)
        else:
            cv2.imwrite(filepath, frame)
        
        print(f"结果已保存至: {filepath}")
        return filepath
    
    def show_comparison(self, original_frame, processed_frame, title1="原始图像", title2="处理结果"):
        """显示原始图像和处理结果的对比
        
        Args:
            original_frame: 原始图像
            processed_frame: 处理后的图像
            title1: 第一张图的标题
            title2: 第二张图的标题
        """
        plt.figure(figsize=(12, 6))
        
        # 显示原始图像
        plt.subplot(121)
        plt.title(title1)
        plt.imshow(original_frame)
        plt.axis('off')
        
        # 显示处理结果
        plt.subplot(122)
        plt.title(title2)
        plt.imshow(processed_frame)
        plt.axis('off')
        
        plt.tight_layout()
        plt.show()
    
    def create_heatmap(self, frame, detections, radius=15):
        """创建目标密度热力图
        
        Args:
            frame: 输入图像
            detections: 检测结果
            radius: 热图点半径
        
        Returns:
            热力图叠加的图像
        """
        if detections is None or len(detections) == 0:
            return frame.copy()
        
        # 创建热力图
        heatmap = np.zeros((frame.shape[0], frame.shape[1]), dtype=np.float32)
        
        # 在检测框中心添加热力点
        for xyxy in detections.xyxy:
            center_x = int((xyxy[0] + xyxy[2]) / 2)
            center_y = int((xyxy[1] + xyxy[3]) / 2)
            
            # 在中心位置创建圆形热区
            cv2.circle(heatmap, (center_x, center_y), radius, 1.0, -1)
        
        # 应用高斯模糊使热力图更平滑
        heatmap = cv2.GaussianBlur(heatmap, (radius*4+1, radius*4+1), 0)
        
        # 归一化热力图
        heatmap = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        
        # 转换为伪彩色
        heatmap_color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        heatmap_color = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB)
        
        # 将热力图叠加到原始图像上
        overlay = cv2.addWeighted(frame, 0.7, heatmap_color, 0.3, 0)
        
        return overlay

# 注意：不创建全局实例，在实际使用时通过导入类并实例化
# 遵循Java风格的实例化模式，在使用时才创建实例